#include <tomcrypt.h>
#include <sys/types.h>
#include <sys/stat.h>

#pragma region variable declaration
int err, hash_idx, prng_idx;
//char *in, *out; // file name
                // file pointer
rsa_key key;

#pragma endregion

#pragma region tool function

int registerall()
{
    /* register prng/hash */
    if (register_prng(&sprng_desc) == -1)
    {
        printf("Error registering sprng");
        return EXIT_FAILURE;
    }
    /* register a math library (in this case TomsFastMath)*/
    ltc_mp = tfm_desc;
    //printf("%s\n", ltc_mp.name);
    if (register_hash(&sha1_desc) == -1)
    {
        printf("Error registering sha1");
        return EXIT_FAILURE;
    }
    hash_idx = find_hash("sha1");
    prng_idx = find_prng("sprng");
}

int usage(char *name)
{
    printf("Usage generate rsa key: %s rsakey uk rk\n", name);
    printf("Usage encrypt: %s encrypt uk plaintext ciphertext\n", name);
    printf("Usage decrypt: %s decrypt rk ciphertext decipher\n", name);
    printf("Usage signature: %s sign rk doc signedfile\n", name);
    printf("Usage verify signature: %s verify uk doc signedfile\n", name);
    printf("Usage generate Diffe-Hellman key: %s dhkey dhuk dhrk\n", name);
    printf("Usage generate secure using dh key: %s dhsecure myrk otheruk secure\n", name);
}

void generatekey(char *uk, char *rk)
{
    unsigned char tmpkey[1024];
    FILE *fdrk, *fduk;

    fduk = fopen(uk, "wb");
    if (fduk == NULL)
    {
        perror("Can't open uk for writing");
        exit(-1);
    }

    fdrk = fopen(rk, "wb");
    if (fdrk == NULL)
    {
        perror("Can't open rk for writing");
        exit(-1);
    }

    /* make an RSA-1024 key */
    if ((err = rsa_make_key(NULL,     /* PRNG state */
                            prng_idx, /* PRNG idx */
                            1024 / 8, /* 1024-bit key */
                            65537,    /* we like e=65537 */
                            &key)     /* where to store the key */
         ) != CRYPT_OK)
    {
        printf("rsa_make_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    unsigned long keylength = 1024;

    if ((err = rsa_export(tmpkey, &keylength, PK_PUBLIC, &key)) != CRYPT_OK)
    {
        printf("rsa_export_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    //printf("%d\n", keylength);

    if (fwrite(tmpkey, 1, keylength, fduk) != keylength)
    {
        printf("Error writing to output.\n");
        exit(-1);
    }

    keylength = 1024;

    if ((err = rsa_export(tmpkey, &keylength, PK_PRIVATE, &key)) != CRYPT_OK)
    {
        printf("rsa_export_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    //printf("%d\n", keylength);

    if (fwrite(tmpkey, 1, keylength, fdrk) != keylength)
    {
        printf("Error writing to output.\n");
        exit(-1);
    }
}

int check_file_size(char *file_name, long long max_size)
{
    struct stat file_info;
    stat(file_name, &file_info);
    long long octets = file_info.st_size;
    if (octets > max_size)
    {
        perror("file oversize");
        exit(-1);
    }
    if (octets < 1)
    {
        perror("empty file");
        exit(-1);
    }
}

void encrypt(char *uk, char *pt, char *ct)
{
    unsigned char plaintext[1024], ciphertext[1024], tmpkey[1024];
    FILE *fduk, *fdpt, *fdct;

    //check_file_size(pt, 512/8);

    fduk = fopen(uk, "rb");
    if (fduk == NULL)
    {
        perror("Can't open uk for reading");
        exit(-1);
    }

    unsigned long keylength = fread(tmpkey, 1, 1024, fduk);
    // printf("%lu\n", keylength);

    if ((err = rsa_import(tmpkey, keylength, &key)) != CRYPT_OK)
    {
        printf("rsa_import %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    fdpt = fopen(pt, "rb");
    if (fdpt == NULL)
    {
        perror("Can't open pt for reading");
        exit(-1);
    }

    unsigned long ptlength = fread(plaintext, 1, 1024, fdpt);
    // printf("%lu\n", ptlength);

    unsigned long ctlength = sizeof(ciphertext);

    if ((err = rsa_encrypt_key(plaintext,  /* data we wish to encrypt */
                               ptlength,   /* data is 16 bytes long */
                               ciphertext, /* where to store ciphertext */
                               &ctlength,  /* length of ciphertext */
                               "TestApp",  /* our lparam for this program */
                               7,          /* lparam is 7 bytes long */
                               NULL,       /* PRNG state */
                               prng_idx,   /* prng idx */
                               hash_idx,   /* hash idx */
                               &key)       /* our RSA key */
         ) != CRYPT_OK)
    {
        printf("rsa_encrypt_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    // printf("%lu\n", ctlength);

    fdct = fopen(ct, "wb");
    if (fdct == NULL)
    {
        perror("Can't open ct for writing");
        exit(-1);
    }

    if (fwrite(ciphertext, 1, ctlength, fdct) != ctlength)
    {
        printf("Error writing ct to output.\n");
        exit(-1);
    }
}

void decrypt(char *rk, char *ct, char *pt)
{
    unsigned char plaintext[1024], ciphertext[1024], tmpkey[1024];
    FILE *fdrk, *fdct, *fdpt;

    fdrk = fopen(rk, "rb");
    if (fdrk == NULL)
    {
        perror("Can't open rk for reading");
        exit(-1);
    }
    unsigned long keylength = fread(tmpkey, 1, 1024, fdrk);
    // printf("%lu\n", keylength);

    if ((err = rsa_import(tmpkey, keylength, &key)) != CRYPT_OK)
    {
        printf("rsa_import %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    //rsa_import(tmpkey, keylength, &key);

    fdct = fopen(ct, "rb");
    if (fdct == NULL)
    {
        perror("Can't open ct for reading");
        exit(-1);
    }
    unsigned long ctlength = fread(ciphertext, 1, 1024, fdct);
    // printf("%lu\n", ctlength);

    unsigned long ptlength = sizeof(plaintext);
    int res;
    if ((err = rsa_decrypt_key(ciphertext, /* encrypted data */
                               ctlength,   /* length of ciphertext */
                               plaintext,  /* where to put plaintext */
                               &ptlength,  /* plaintext length */
                               "TestApp",  /* lparam for this program */
                               7,          /* lparam is 7 bytes long */
                               hash_idx,   /* hash idx */
                               &res,       /* validity of data */
                               &key)       /* our RSA key */
         ) != CRYPT_OK)
    {
        printf("rsa_decrypt_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    if (!res)
    {
        printf("%s\n", "cipher text invalid");
        exit(EXIT_FAILURE);
    }
    // printf("%s\n", plaintext);

    fdpt = fopen(pt, "wb");
    if (fdpt == NULL)
    {
        perror("Can't open pt for writing");
        exit(-1);
    }

    if (fwrite(plaintext, 1, ptlength, fdpt) != ptlength)
    {
        printf("Error writing pt to output.\n");
        exit(-1);
    }
}

void sign(char *rk, char *pt, char *sg)
{
    unsigned char hashbuf[1024], sgbuf[1024], tmpkey[1024];
    FILE *fdrk, *fdsg;

    fdrk = fopen(rk, "rb");
    if (fdrk == NULL)
    {
        perror("Can't open rk for reading");
        exit(-1);
    }
    unsigned long keylength = fread(tmpkey, 1, 1024, fdrk);
    // printf("%lu\n", keylength);

    if ((err = rsa_import(tmpkey, keylength, &key)) != CRYPT_OK)
    {
        printf("rsa_import %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    unsigned long hashlength = sizeof(hashbuf);
    if ((err = hash_file(hash_idx, pt, hashbuf, &hashlength)) != CRYPT_OK)
    {
        printf("hash_file %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    // printf("%lu\n", hashlength);

    unsigned long sglength = sizeof(sgbuf);
    int saltlen = rsa_sign_saltlen_get_max(hash_idx, &key);

    if ((err = rsa_sign_hash(hashbuf,
                             hashlength,
                             sgbuf,
                             &sglength,
                             NULL,
                             prng_idx,
                             hash_idx,
                             saltlen,
                             &key)) != CRYPT_OK)
    {
        printf("rsa_import %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    fdsg = fopen(sg, "wb");
    if (fdsg == NULL)
    {
        perror("Can't open sg for writing");
        exit(-1);
    }

    if (fwrite(sgbuf, 1, sglength, fdsg) != sglength)
    {
        printf("Error writing pt to output.\n");
        exit(-1);
    }
}

void verify(char *uk, char *pt, char *sg)
{
    unsigned char hashbuf[1024], sgbuf[1024], tmpkey[1024];
    FILE *fduk, *fdsg;

    fduk = fopen(uk, "rb");
    if (fduk == NULL)
    {
        perror("Can't open uk for reading");
        exit(-1);
    }
    unsigned long keylength = fread(tmpkey, 1, 1024, fduk);
    // printf("key length : %lu\n", keylength);

    if ((err = rsa_import(tmpkey, keylength, &key)) != CRYPT_OK)
    {
        printf("rsa_import %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    unsigned long hashlength = sizeof(hashbuf);
    if ((err = hash_file(hash_idx, pt, hashbuf, &hashlength)) != CRYPT_OK)
    {
        printf("hash_file %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    // printf("hash length : %lu\n", hashlength);

    fdsg = fopen(sg, "rb");
    if (fdsg == NULL)
    {
        perror("Can't open sg for reading");
        exit(-1);
    }
    unsigned long sglength = fread(sgbuf, 1, 1024, fdsg);
    // printf("signature length : %lu\n", sglength);

    int res;
    int saltlen = rsa_sign_saltlen_get_max(hash_idx, &key);

    if ((err = rsa_verify_hash(
             sgbuf,
             sglength,
             hashbuf,
             hashlength,
             hash_idx,
             saltlen,
             &res,
             &key)) != CRYPT_OK)
    {
        printf("rsa_verify %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    if (res == 1)
    {
        printf("%s\n", "succedd on verify");
    }
    else
    {
        printf("%s\n", "fail on verify");
    }
}

int dh(char *uk, char *rk)
{
    unsigned char ukbuf[4096], rkbuf[4096];
    FILE *fduk, *fdrk;
    dh_key dhkey;

    if ((err = dh_set_pg_groupsize(128, &dhkey)) != CRYPT_OK)
    {
        printf("dh_set_pg_groupsize %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    if ((err = dh_generate_key(NULL, prng_idx, &dhkey)) != CRYPT_OK)
    {
        printf("dh_generate_key %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    unsigned long uklength = sizeof(ukbuf);
    if ((err = dh_export(ukbuf, &uklength, PK_PUBLIC, &dhkey)) != CRYPT_OK)
    {
        printf("dh_export uk %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    // printf("uk length : %lu\n", uklength);

    unsigned long rklength = sizeof(rkbuf);
    if ((err = dh_export(rkbuf, &rklength, PK_PRIVATE, &dhkey)) != CRYPT_OK)
    {
        printf("dh_export rk %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }
    // printf("rk length : %lu\n", rklength);

    fduk = fopen(uk, "wb");
    if (fduk == NULL)
    {
        perror("Can't open uk for writing");
        exit(-1);
    }
    if (fwrite(ukbuf, 1, uklength, fduk) != uklength)
    {
        printf("Error writing uk to output.\n");
        exit(-1);
    }

    fdrk = fopen(rk, "wb");
    if (fdrk == NULL)
    {
        perror("Can't open rk for writing");
        exit(-1);
    }
    if (fwrite(rkbuf, 1, rklength, fdrk) != rklength)
    {
        printf("Error writing rk to output.\n");
        exit(-1);
    }
}

void secure(char *myrk, char *otheruk, char *secure)
{
    unsigned char mrkbuf[4096], oukbuf[4096], sebuf[4096];
    FILE *fdmrk, *fdouk, *fdse;
    dh_key mrkey, oukey;

    fdmrk = fopen(myrk, "rb");
    if (fdmrk == NULL)
    {
        perror("Can't open my rk for reading");
        exit(-1);
    }
    unsigned long mrklength = fread(mrkbuf, 1, sizeof(mrkbuf), fdmrk);
    // printf("my rk length : %lu\n", mrklength);

    if ((err = dh_import(mrkbuf, mrklength, &mrkey)) != CRYPT_OK)
    {
        printf("dh_import my rk %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    fdouk = fopen(otheruk, "rb");
    if (fdouk == NULL)
    {
        perror("Can't open other uk for reading");
        exit(-1);
    }
    unsigned long ouklength = fread(oukbuf, 1, sizeof(oukbuf), fdouk);
    // printf("other uk length : %lu\n", ouklength);

    if ((err = dh_import(oukbuf, ouklength, &oukey)) != CRYPT_OK)
    {
        printf("dh_import other uk %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    unsigned long selength = sizeof(sebuf);
    if ((err = dh_shared_secret(&mrkey, &oukey, sebuf, &selength)) != CRYPT_OK)
    {
        printf("dh_shared_secret %s", error_to_string(err));
        exit(EXIT_FAILURE);
    }

    fdse = fopen(secure, "wb");
    if (fdse == NULL)
    {
        perror("Can't open secure for writing");
        exit(-1);
    }
    if (fwrite(sebuf, 1, selength, fdse) != selength)
    {
        printf("Error writing secure to output.\n");
        exit(-1);
    }

    printf("secure after\n");
}

#pragma endregion

int main(int argc, char *argv[])
{
    if (argc < 4)
    {
        return usage(argv[0]);
    }

    printf("%s\n", argv[1]);

    registerall();

    if (!strcmp(argv[1], "rsakey"))
    {
        if (argc != 4)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        generatekey(argv[2], argv[3]);
        return 0;
    }
    if (!strcmp(argv[1], "encrypt"))
    {
        if (argc != 5)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        encrypt(argv[2], argv[3], argv[4]);
        return 0;
    }
    if (!strcmp(argv[1], "decrypt"))
    {
        if (argc != 5)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        decrypt(argv[2], argv[3], argv[4]);
        return 0;
    }
    if (!strcmp(argv[1], "sign"))
    {
        if (argc != 5)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        sign(argv[2], argv[3], argv[4]);
        return 0;
    }
    if (!strcmp(argv[1], "verify"))
    {
        if (argc != 5)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        verify(argv[2], argv[3], argv[4]);
        return 0;
    }
    if (!strcmp(argv[1], "dhkey"))
    {
        if (argc != 4)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        dh(argv[2], argv[3]);
        return 0;
    }
    if (!strcmp(argv[1], "dhsecure"))
    {
        if (argc != 5)
        {
            printf("invalid argument amount\n");
            return usage(argv[0]);
        }
        secure(argv[2], argv[3], argv[4]);
        return 0;
    }

    printf("invalid parameter %s.\n", argv[1]);
    return usage(argv[0]);

}
